"""
多项式拟合示例脚本 - 命令行版本
展示如何在Python代码中使用NumPy进行多项式拟合
"""

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import sys
import argparse

def fit_polynomial(x, y, degree=2):
    """
    使用给定的x和y数据拟合多项式
    
    参数:
        x (array-like): 自变量数据
        y (array-like): 因变量数据
        degree (int): 多项式次数
        
    返回:
        tuple: (多项式函数, 系数数组, 统计信息字典)
    """
    # 拟合多项式
    coeffs = np.polyfit(x, y, degree)
    poly_func = np.poly1d(coeffs)
    
    # 计算预测值
    y_pred = poly_func(x)
    
    # 计算统计指标
    mse = np.mean((y - y_pred) ** 2)
    rmse = np.sqrt(mse)
    
    # 计算R²
    y_mean = np.mean(y)
    ss_total = np.sum((y - y_mean) ** 2)
    ss_residual = np.sum((y - y_pred) ** 2)
    r2 = 1 - (ss_residual / ss_total)
    
    # 返回结果
    stats = {
        'mse': mse,
        'rmse': rmse,
        'r2': r2
    }
    
    return poly_func, coeffs, stats

def format_polynomial(coeffs, degree):
    """
    格式化多项式系数为字符串表达式
    
    参数:
        coeffs (array): 多项式系数
        degree (int): 多项式次数
        
    返回:
        str: 格式化的多项式表达式
    """
    formula = "y = "
    for i, coef in enumerate(coeffs):
        power = degree - i
        if i > 0:
            formula += " + " if coef >= 0 else " - "
            formula += f"{abs(coef):.6g}"
        else:
            formula += f"{coef:.6g}"
        
        if power > 0:
            formula += f"x"
            if power > 1:
                formula += f"^{power}"
    
    return formula

def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='多项式拟合示例')
    parser.add_argument('--file', '-f', type=str, help='输入数据文件 (CSV或Excel)')
    parser.add_argument('--degree', '-d', type=int, default=2, help='多项式次数')
    parser.add_argument('--visualize', '-v', action='store_true', help='显示可视化图表')
    args = parser.parse_args()
    
    # 使用样例数据或读取文件
    if args.file:
        try:
            if args.file.endswith(('.xlsx', '.xls')):
                data = pd.read_excel(args.file)
            elif args.file.endswith('.csv'):
                data = pd.read_csv(args.file)
            else:
                raise ValueError("不支持的文件格式")
            
            x = data.iloc[:, 0].values
            y = data.iloc[:, 1].values
            
        except Exception as e:
            print(f"错误: {e}")
            return
    else:
        # 使用样例数据
        print("使用样例数据")
        x = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
        y = np.array([2.1, 4.3, 8.7, 16.2, 25.9, 36.2, 49.1, 63.8, 81.2, 100.5])
    
    # 拟合多项式
    poly_func, coeffs, stats = fit_polynomial(x, y, args.degree)
    
    # 显示结果
    print("\n=== 多项式拟合结果 ===")
    print(f"数据点数量: {len(x)}")
    print(f"多项式次数: {args.degree}")
    print(f"\n多项式公式:\n{format_polynomial(coeffs, args.degree)}")
    print(f"\n统计信息:")
    print(f"  均方误差 (MSE): {stats['mse']:.6g}")
    print(f"  均方根误差 (RMSE): {stats['rmse']:.6g}")
    print(f"  确定系数 (R²): {stats['r2']:.6g}")
    
    # 预测示例
    print("\n预测示例:")
    for x_val in [min(x), np.mean(x), max(x)]:
        y_pred = poly_func(x_val)
        print(f"  x = {x_val:.4g} → 预测 y = {y_pred:.6g}")
    
    # 可视化
    if args.visualize:
        plt.figure(figsize=(10, 6))
        
        # 绘制原始数据点
        plt.scatter(x, y, color='blue', label='原始数据')
        
        # 创建平滑曲线
        x_line = np.linspace(min(x), max(x), 1000)
        y_line = poly_func(x_line)
        
        # 绘制拟合曲线
        plt.plot(x_line, y_line, 'r-', label='拟合曲线')
        
        # 计算并显示预测值
        y_pred = poly_func(x)
        
        # 绘制预测点
        plt.scatter(x, y_pred, color='green', alpha=0.5, label='预测值')
        
        # 绘制从实际到预测的线
        for i in range(len(x)):
            plt.plot([x[i], x[i]], [y[i], y_pred[i]], 'gray', linestyle='--', alpha=0.3)
        
        plt.title('多项式拟合')
        plt.xlabel('X')
        plt.ylabel('Y')
        plt.grid(True)
        plt.legend()
        plt.tight_layout()
        plt.show()

if __name__ == "__main__":
    main() 